struct Solution;

impl Solution {
    pub fn matrix_block_sum(mat: Vec<Vec<i32>>, k: i32) -> Vec<Vec<i32>> {
        unsafe fn get_num(src: &[Vec<i32>], m: usize, n: usize) -> i32 {
            *src.get_unchecked(m).get_unchecked(n)
        }

        unsafe fn get_num_mut(src: &mut Vec<Vec<i32>>, m: usize, n: usize) -> &mut i32 {
            src.get_unchecked_mut(m).get_unchecked_mut(n)
        }

        // 长度
        let (m, n) = (mat.len(), mat[0].len());
        let mut sum = vec![vec![0; n + 1]; m + 1];

        for i in 0..m {
            for j in 0..n {
                unsafe {
                    let tmp = get_num(&sum, i + 1, j) +
                        get_num(&sum, i, j + 1) -
                        get_num(&sum, i, j) +
                        get_num(&mat, i, j);
                    *get_num_mut(&mut sum, i + 1, j + 1) = tmp;
                }
            }
        }

        let mut res = vec![vec![0; n]; m];
        let k = k as usize;
        for (i, vec) in res.iter_mut().enumerate() {
            for (j, val) in vec.iter_mut().enumerate() {
                let left = if j > k {
                    j - k
                } else {
                    0
                };
                let right = std::cmp::min(j + k + 1, n);
                let top = if i > k {
                    i - k
                } else {
                    0
                };
                let down = std::cmp::min(i + k + 1, m);
                unsafe {
                    *val = get_num(&sum, down, right) -
                        get_num(&sum, top, right) -
                        get_num(&sum, down, left) +
                        get_num(&sum, top, left);
                }
            }
        }


        res
    }
}

#[cfg(test)]
mod test {
    use crate::Solution;

    #[test]
    fn test_01() {
        let mat = vec![vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9]];
        let res = Solution::matrix_block_sum(mat, 1);
        assert_eq!(res, vec![vec![12, 21, 16], vec![27, 45, 33], vec![24, 39, 28]]);
    }

    #[test]
    fn test_02() {
        let mat = vec![vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9]];
        let res = Solution::matrix_block_sum(mat, 2);
        assert_eq!(res, vec![vec![45, 45, 45], vec![45, 45, 45], vec![45, 45, 45]]);
    }
}


fn main() {
    Solution::matrix_block_sum(vec![vec![]], 1);
}
